801bc5
@@ -27,9 +27,6 @@
import static org.wildfly.extension.messaging.activemq.Capabilities.ACTIVEMQ_SER
 import org.jboss.as.controller.AbstractRemoveStepHandler;
 import org.jboss.as.controller.OperationContext;
 import org.jboss.as.controller.OperationFailedException;
-import org.jboss.as.controller.OperationStepHandler;
-import org.jboss.as.controller.PathAddress;
-import org.jboss.as.controller.registry.Resource;
 import org.jboss.dmr.ModelNode;
 import org.jboss.msc.service.ServiceName;
 import org.wildfly.extension.messaging.activemq.jms.JMSServices;
@@ -49,27 +46,6 @@
class ServerRemove extends AbstractRemoveStepHandler {
     }
 
 
-    @Override
-    protected void performRemove(OperationContext context, ModelNode operation, ModelNode model) throws OperationFailedException {
-        Resource resource = context.readResource(PathAddress.EMPTY_ADDRESS);
-        // add a runtime step to remove services related to broadcast-group/discovery-group that are started
-        // when the server is added.
-        context.addStep(new OperationStepHandler() {
-            @Override
-            public void execute(OperationContext context, ModelNode operation) throws OperationFailedException {
-                final String serverName = context.getCurrentAddressValue();
-                final ServiceName serviceName = MessagingServices.getActiveMQServiceName(serverName);
-                for(final Resource.ResourceEntry broadcastGroup : resource.getChildren(CommonAttributes.BROADCAST_GROUP)) {
-                    context.removeService(GroupBindingService.getBroadcastBaseServiceName(serviceName).append(broadcastGroup.getName()));
-                }
-                for(final Resource.ResourceEntry divertGroup : resource.getChildren(CommonAttributes.DISCOVERY_GROUP)) {
-                    context.removeService(GroupBindingService.getDiscoveryBaseServiceName(serviceName).append(divertGroup.getName()));
-                }
-            }
-        }, OperationContext.Stage.RUNTIME);
-        super.performRemove(context, operation, model);
-    }
-
     @Override
     protected void performRuntime(OperationContext context, ModelNode operation, ModelNode model) throws OperationFailedException {
 
@@ -77,5 +53,17 @@
class ServerRemove extends AbstractRemoveStepHandler {
         final ServiceName serviceName = MessagingServices.getActiveMQServiceName(serverName);
         context.removeService(JMSServices.getJmsManagerBaseServiceName(serviceName));
         context.removeService(MessagingServices.getActiveMQServiceName(serverName));
+        // remove services related to broadcast-group/discovery-group that are started
+        // when the server is added
+        if (model.hasDefined(CommonAttributes.BROADCAST_GROUP)) {
+            for (String name : model.get(CommonAttributes.BROADCAST_GROUP).keys()) {
+                context.removeService(GroupBindingService.getBroadcastBaseServiceName(serviceName).append(name));
+            }
+        }
+        if (model.hasDefined(CommonAttributes.DISCOVERY_GROUP)) {
+            for (String name : model.get(CommonAttributes.DISCOVERY_GROUP).keys()) {
+                context.removeService(GroupBindingService.getDiscoveryBaseServiceName(serviceName).append(name));
+            }
+        }
     }
 }
